# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import json
import logging
import typing

import requests
import tornado.gen

import coin.flow.subscriber as flow_subscriber
import coin.flow.topic_map as topic_map

from coin.base.config import Config
from coin.exchange.base.fill_checker_order_info import from_proto_order_side
from coin.exchange.base.order_fill_util import FillManager, LastFill
from coin.exchange.base.rest.task_scheduler import Scheduler, wrap_async_task
from coin.exchange.korbit_v1.kr_rest.constants import min_order_qty
from coin.exchange.korbit_v1.kr_rest.private_client import KorbitPrivateParser
from coin.exchange.korbit_v1.kr_rest.product import KorbitProduct, KorbitCurrency
from coin.exchange.util.rate_limit import RateLimiter
from coin.exchange.util.rest_rate_limit import RestRateLimiter
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig, OrderEvent

from coin.exchange.korbit_v1.kr_rest.native_private_client import (
    KorbitNativePrivateClient,
    KorbitOrderType,
    KorbitOrderStatus,
    KorbitCancelOrderStatus,
    KorbitPlaceOrderStatus,
    KorbitAuthGrantType,
)
from coin.exchange.base.rest_order_gateway_base import (
    BaseBalanceOrderGatewayInfo,
    RestOrderGateway,
    PrivateExchangeMessage,
    exchange_message_handler,
    create_order_internal,
    Order,
    OrderSide,
    OrderType,
    ProductFillElement,
    to_proto_order_side,
)


def create_order_from_proto(proto_order, timestamp, order_id=None):
  price_orig = proto_order.price
  qty_orig = proto_order.qty
  order_side = from_proto_order_side(proto_order.side)
  order_type = OrderType.LIMIT
  product = KorbitProduct.FromStrProduct(proto_order.symbol)
  exchange_order_id = _gen_unique_exchange_order_id(product, proto_order.order_id)

  fill = FillManager(
      qty_orig=qty_orig,
      price_orig=price_orig,
      qty_accum_fill=0.0,
      price_avg_fill=None,
      qty_last_fill=None,
      price_last_fill=None,
  )

  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=None,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=timestamp,
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                ))
  return order


def _normalize_not_enough_status(status_str):
  if status_str.startswith('not_enough'):
    return 'not_enough'
  else:
    return status_str


def _to_normalized_order_side(korbit_order_side_str: str):
  if korbit_order_side_str == 'ask':
    return OrderSide.SELL
  elif korbit_order_side_str == 'bid':
    return OrderSide.BUY
  else:
    raise ValueError('Unknown korbit order side: %s', korbit_order_side_str)


def _gen_unique_exchange_order_id(product, product_exchange_order_id):
  return '%s.%s' % (product.symbol, product_exchange_order_id)


def _get_product_exchange_order_id(exchange_order_id):
  res = exchange_order_id.split('.')
  assert len(res) == 2, exchange_order_id
  return res[1]


def create_order_from_order_msg(order_msg, timestamp):
  """
  {
    "id": "89999",
    "currency_pair": "btc_krw",
    "side": "bid",
    "avg_price": "2900000",
    "price": "3000000",
    "order_amount": "0.81140000",
    "filled_amount": "0.33122200",
    "order_total": "2434200",
    "filled_total": "960543",
    "created_at": "1500033942638",
    "last_filled_at": "1500533946947",
    "status": "partially_filled",
    "fee": "0.00000500"
  },
  """
  order_id = None
  product = KorbitProduct.FromStrNativeProduct(order_msg['currency_pair'])
  exchange_order_id = _gen_unique_exchange_order_id(product, order_msg['id'])
  order_side = _to_normalized_order_side(order_msg['side'])
  last_filled_at = order_msg.get('last_filled_at')
  fill_id = '%s_%s' % (exchange_order_id, last_filled_at)

  if 'price' in order_msg and 'order_amount' in order_msg:
    price_orig = float(order_msg['price'])
    qty_orig = float(order_msg['order_amount'])
  else:
    logging.info('No price/qty found, maybe market order!\n%s', json.dumps(order_msg, indent=2))
    return None

  price_avg_fill = None
  if 'avg_price' in order_msg:
    price_avg_fill = float(order_msg['avg_price'])
  else:
    logging.debug('avg price not found!\n%s', json.dumps(order_msg, indent=2))

  qty_accum_fill = float(order_msg['filled_amount'])
  order_type = OrderType.LIMIT
  status = KorbitOrderStatus(order_msg['status'])
  fee = float(order_msg.get('fee', 0))

  fill = FillManager(
      price_orig=price_orig,
      qty_orig=qty_orig,
      price_last_fill=price_avg_fill,
      qty_last_fill=qty_accum_fill,
      price_avg_fill=price_avg_fill,
      qty_accum_fill=qty_accum_fill,
      timestamp_last_fill=timestamp,
  )
  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=fill_id,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=timestamp,
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                    status=status,
                    fee=fee,
                ))
  return order


class KorbitOrderGatewayInfo(BaseBalanceOrderGatewayInfo):
  CurrencyType = KorbitCurrency
  ProductType = KorbitProduct

  def update_balance(self, update_msg):
    account_balance = KorbitPrivateParser.parse_balance(update_msg)
    for balance in account_balance.each_balance:
      currency = KorbitCurrency.FromStrCurrency(balance.currency)
      if currency in self.currencies:
        self.set_balance(currency, balance.total, balance)


class KorbitOrderGateway(RestOrderGateway):
  _exchange = 'Korbit'
  _market_type = 'Spot'
  _api_version = None

  OGInfoType = KorbitOrderGatewayInfo

  class _MsgType(enum.IntEnum):
    REST_ACCOUNT_BALANCE = 1
    REST_OPEN_ORDERS = 2
    REST_USER_ORDERS = 3
    REST_ORDER_FILLS = 4
    AUTH_TOKEN = 5

  def __init__(self,
               currencies: typing.List[KorbitCurrency],
               products: typing.List[KorbitProduct],
               config: Config,
               og_config=None,
               logger=None):
    super().__init__(currencies, products, config, og_config, logger)
    self._private_client = None
    self._scheduler = None
    self._ioloop = None
    self._flow_sub = None
    self._publisher = None

    default_og_config = OrderGatewayConfig(
        account_query_period=2,
        orders_snapshot_query_period=10,
        fill_query_period=2,
        hit_rate_limit_cooldown_time=12,
    )
    default_og_config.MergeFrom(self._og_config)
    self._og_config = default_og_config

    # Handle fill message
    self._rest_rate_limiter = RestRateLimiter(delay=5)
    self._order_rate_limiter = RateLimiter(
        self._og_config.order_rate_limit,
        self._og_config.order_rate_limit_per,
    )
    self._first_query_map = set()

  def start(self, queue, ioloop, flow_sub=None):
    api_host = self._og_config.rest_root_url \
               if self._og_config.HasField('rest_root_url') else None
    self._private_client = KorbitNativePrivateClient(
        key_file=self._config.key_file,
        max_connections=self._og_config.max_connections,
        api_host=api_host,
        use_async=True,
        rest_request_gap_sec=self._og_config.rest_request_gap_sec,
    )
    topic_id = topic_map.search_topic_id('korbit_rest_private_1')[0][1]
    self._ioloop = ioloop
    self._publisher = queue.get_writer(topic_id=topic_id)
    self._scheduler = Scheduler(ioloop=ioloop, logger=self._logger)
    self._install_periodic_query()

    # Subscribe to flow for feed.
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe('korbit_rest_private_1', self.on_private_exchange_msg)
    self._logger.info(
        'order_rate_limit: %s, order_rate_limit_per: %s',
        self._og_config.order_rate_limit,
        self._og_config.order_rate_limit_per,
    )

  def submittable(self):
    return self.is_ready() \
           and self._rest_rate_limiter.is_available() \
           and self._order_rate_limiter.check()

  @tornado.gen.coroutine
  def _install_periodic_query(self):
    assert self._scheduler is not None, 'self._scheduler is not set!'
    assert self._publisher is not None, 'self._publisher is not set!'

    query_token_task = wrap_async_task(
        self._private_client.query_token,
        msg_type=self._MsgType.AUTH_TOKEN,
        publisher=self._publisher,
        func_params=dict(grant_type=KorbitAuthGrantType.PASSWORD),
    )

    for _ in range(3):
      if self._private_client.is_ready():
        break

      self._scheduler.run_once(query_token_task)
      for _ in range(10):
        if self._private_client.is_ready():
          break
        else:
          yield tornado.gen.sleep(1)
          self._logger.info('Wait until korbit auth token is ready!')

    if not self._private_client.is_ready():
      self._logger.error('Failed to get auth token!!')
      return

    # query balance
    query_balance_task = wrap_async_task(
        self._private_client.query_user_balances,
        msg_type=self._MsgType.REST_ACCOUNT_BALANCE,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(period=self._og_config.account_query_period,
                                      job_func=query_balance_task)

    # query active orders
    for product in self._og_info.products:
      yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
      query_orders_task = wrap_async_task(
          self._private_client.query_open_orders,
          msg_type=self._MsgType.REST_OPEN_ORDERS,
          publisher=self._publisher,
          func_params=dict(currency_pair=product.native_symbol),
      )
      self._scheduler.add_periodic_task(
          period=self._og_config.orders_snapshot_query_period,
          job_func=query_orders_task,
      )

      yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
      query_orders_task = wrap_async_task(
          self._private_client.query_user_orders,
          msg_type=self._MsgType.REST_USER_ORDERS,
          publisher=self._publisher,
          func_params=dict(currency_pair=product.native_symbol),
      )
      self._scheduler.add_periodic_task(
          period=self._og_config.fill_query_period,
          job_func=query_orders_task,
      )

    self._scheduler.add_periodic_task(
        period=1,
        job_func=self.refresh_token,
        immediate_first_run=False,
    )
    self.start_scheduler()

  def refresh_token(self):
    if self._private_client.need_refresh_token(tolerance_time_sec=60):
      refresh_token_task = wrap_async_task(
          self._private_client.query_token,
          msg_type=self._MsgType.AUTH_TOKEN,
          publisher=self._publisher,
          func_params=dict(grant_type=KorbitAuthGrantType.REFRESH_TOKEN),
      )
      self._scheduler.run_once(refresh_token_task)

  @tornado.gen.coroutine
  def _on_hit_order_rate_limit(self):
    if not self.is_ready():
      return

    self._og_ready.set_reinit()
    self._scheduler.stop()
    self._logger.info('Hit rate limit, disable OG for % seconds.',
                      self._og_config.hit_rate_limit_cooldown_time)
    yield tornado.gen.sleep(self._og_config.hit_rate_limit_cooldown_time)
    self._scheduler.start()
    self._logger.info('Korbit OG is active again.')

  def _async_submit_order(self, order, post_only=False):
    if post_only:
      self._logger.debug('post only not supported.')
    assert order.order_id is not None, order
    assert order.order_type == OrderType.LIMIT, order

    min_qty = min_order_qty[order.product.base.currency]
    if order.qty < min_qty:
      self._logger.error('order qty is less than min order qty! %s, %s', order.product, order.qty)
      self._order_manager.remove_order(order)
      return

    formatter = self._order_formatter[order.product]
    params = {
        'currency_pair': order.product.native_symbol,
        'order_type': KorbitOrderType.LIMIT,
        'price': formatter.format_price(order.price),
        'coin_amount': formatter.format_qty(order.qty),
    }
    if order.order_side == OrderSide.BUY:
      return self._private_client.place_buy_order(**params)
    elif order.order_side == OrderSide.SELL:
      return self._private_client.place_sell_order(**params)
    else:
      raise ValueError('Invalid order side : %s', order.order_side)

  def _async_cancel_order(self, order):
    if order.internal.exchange_order_id is None:
      self._logger.error('Order not accepted yet! %s', order.order_id)
      return

    exchange_order_id = order.internal.exchange_order_id
    fut = self._private_client.cancel_orders(
        currency_pair=order.product.native_symbol,
        order_id_list=[_get_product_exchange_order_id(exchange_order_id)],
    )
    return fut

  def _handle_error_response(self, response: requests.Response):
    if response.status_code == 504:
      self._logger.error('504 Gateway Timeout error!')
    elif response.status_code == 429:
      self._rest_rate_limiter.set_error(response.status_code)
      self._logger.error('Rate limit hit (429)!')

    self._logger.error(
        'response error! %s, %s, %s, %s',
        response.status_code,
        response.content,
        response.url,
        response.request.body,
    )
    self.log_json_message(dict(response.headers), msg='Response header')

  def _handle_submit_response(self, order, response):
    """
    {
      "orderId": 23733520,
      "status":"success",
      "currencyPair": "xrp_krw"
    }
    """
    if response.status_code != 200:
      self._handle_error_response(response)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_REJECTED

    msg = response.json()
    status_str = _normalize_not_enough_status(msg['status'])
    exchange_order_id = _gen_unique_exchange_order_id(order.product, msg['orderId'])
    status = KorbitPlaceOrderStatus(status_str)
    if status == KorbitPlaceOrderStatus.SUCCESS:
      prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
      if prev_order is None:
        self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
      else:
        self._logger.info('order id replaced: %s => %s', prev_order.order_id, order.order_id)
        self._order_manager.remove_order(order)
        self._order_manager.remove_order(prev_order)
        prev_order.order_id = order.order_id
        self._order_manager.add_order(prev_order)
      return OrderEvent.ORDER_ACCEPTED
    elif status in (KorbitPlaceOrderStatus.NOT_ENOUGH, KorbitPlaceOrderStatus.TOO_MANY_ORDERS):
      self._order_manager.remove_order(order, ignore_error=True)
      self._logger.error('Submit order Error! %s', msg)
      return OrderEvent.ORDER_REJECTED
    else:
      self._order_manager.remove_order(order, ignore_error=True)
      self._logger.error('Submit order Error! %s', msg)
      return OrderEvent.ORDER_ERROR

  def _handle_cancel_response(self, order, response):
    """ Cancel response has a very simple form:
    [
      {"orderId":"1001","status":"not_found"},
    ]
    """
    if response.status_code != 200:
      self._handle_error_response(response)
      order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR

    msg = response.json()
    if len(msg) > 1:
      self._logger.error('Weird cancel response! %s', msg)

    msg = msg[0]
    exchange_order_id = _gen_unique_exchange_order_id(order.product, msg['orderId'])
    if order.internal.exchange_order_id != exchange_order_id:
      self._logger.error('Insane cancel response: %s, %s', order, msg)
      return OrderEvent.CANCEL_ERROR

    status = KorbitCancelOrderStatus(msg['status'])
    if status == KorbitCancelOrderStatus.SUCCESS:
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, cancel_confirmed=True)
      else:
        self._logger.info('Order %s is dead already!', order.order_id)
      return OrderEvent.CANCEL_CONFIRMED
    elif status in (KorbitCancelOrderStatus.NOT_FOUND,
                    KorbitCancelOrderStatus.ALREADY_FILLED,
                    KorbitCancelOrderStatus.ALREADY_CANCELED):
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, cancel_confirmed=True)
      else:
        self._logger.info('Order %s is dead already!', order.order_id)
      return OrderEvent.CANCEL_CONFIRMED
    else:
      order.cancel_sent = False
      self._logger.error('Order is dead! %s', order.order_id)
      return OrderEvent.CANCEL_ERROR

  def _update_order(self, order_info, timestamp):
    exchange_order_id = order_info.internal.exchange_order_id
    fill_id = order_info.tag
    order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if order is None:
      order_info.order_id = self._order_id_generator.gen_order_id(is_foreign=True)
      self._order_manager.add_order(order_info)
      order = order_info

    if order.accepted is False:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)

    price_avg_fill = order_info.internal.fill.value.price_avg_fill
    qty_accum_fill = order_info.internal.fill.value.qty_accum_fill
    last_fill = None

    if price_avg_fill is None:
      return

    if order is not order_info:
      last_fill = order.internal.fill.update_by_accum_fill(price_avg_fill, qty_accum_fill)
    else:
      # if order is a filled/partially filled foreign order,
      # publish its filled part
      if fill_id not in self._processed_fills:
        last_fill = LastFill(timestamp_last_fill=timestamp,
                             price_last_fill=price_avg_fill,
                             qty_last_fill=qty_accum_fill)

    if last_fill:
      if self._order_manager.is_dead_order(order):
        self._logger.error('Weird! Dead order has fill message! %s', order)

      fill_type = self._guess_fill_type(order, last_fill.price_last_fill, timestamp)
      fill_proto = ProductFillElement(price=last_fill.price_last_fill,
                                      qty=last_fill.qty_last_fill,
                                      symbol=order.product.symbol,
                                      order_id=order.order_id,
                                      side=to_proto_order_side(order.order_side),
                                      guess_fill_type=fill_type,
                                      fill_id=fill_id,
                                      exchange_order_id=order.internal.exchange_order_id)
      self.process_fill(fill_proto, timestamp)

    status = order_info.internal.status
    if status == KorbitOrderStatus.FILLED:
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, fully_filled=True)

  @exchange_message_handler(_MsgType.REST_ACCOUNT_BALANCE)
  def _handle_account_balance(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != 200:
      self._handle_error_response(response)
      return

    msg = response.json()
    self._og_info.update_balance(msg)
    self._og_logger.write_balance(self._og_info.balance_as_proto())

    if not self.is_ready():
      self._og_ready.set_balance_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_OPEN_ORDERS)
  def _handle_open_orders(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    timestamp = exchange_message.received_timestamp
    if response.status_code != 200:
      self._handle_error_response(response)
      return

    msg = response.json()
    if not msg:
      return

    currency_pair = exchange_message.request_params['currency_pair']
    product = KorbitProduct.FromStrNativeProduct(currency_pair)
    proto_orders = KorbitPrivateParser.parse_orders(msg, product)
    exchange_order_dict = {}
    products = set()
    for proto_order in proto_orders.each_order:
      order = create_order_from_proto(proto_order, timestamp)
      exchange_order_id = order.internal.exchange_order_id
      prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
      if prev_order is None:
        order_id = self._order_id_generator.gen_order_id(is_foreign=True)
        order.order_id = order_id
      else:
        order.order_id = prev_order.order_id
      if order.product in self._og_info.products:
        products.add(order.product)
        exchange_order_dict[exchange_order_id] = order
    self._order_manager.copy_from_exchange_open_orders(exchange_order_dict, products)

  @exchange_message_handler(_MsgType.REST_USER_ORDERS)
  def _handle_user_orders(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    timestamp = exchange_message.received_timestamp
    if response.status_code != 200:
      self._handle_error_response(response)
      return

    native_symbol = exchange_message.request_params['currency_pair']
    msg = response.json()
    is_first_fill_query = native_symbol not in self._first_query_map
    if is_first_fill_query:
      self._first_query_map.add(native_symbol)

    for order_msg in msg:
      order_info = create_order_from_order_msg(order_msg, timestamp)
      if order_info is None:
        continue

      if order_info.product not in self._og_info.products:
        continue

      if is_first_fill_query:
        if order_info.internal.status in (KorbitOrderStatus.FILLED,
                                          KorbitOrderStatus.PARTIALLY_FILLED):
          fill_id = order_info.tag
          self._processed_fills.add(fill_id)
      self._update_order(order_info, timestamp)

    if not self.is_ready():
      self._og_ready.set_orders_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.AUTH_TOKEN)
  def _handle_refresh_token(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    timestamp = exchange_message.received_timestamp
    if response.status_code != 200:
      self._handle_error_response(response)
      return

    self._logger.info('Update auth token!')
    token_msg = response.json()
    token_fetched_time_sec = timestamp / 1e9
    self._private_client.set_token(token_msg, token_fetched_time_sec)
